package rsaencrypt

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"fmt"
)

// GenerateKeyPair 生成 RSA 密钥对
func GenerateKeyPair(bits int) (*rsa.PrivateKey, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return nil, fmt.Errorf("无法生成密钥对: %v", err)
	}
	return privateKey, nil
}

// EncryptRSA 使用公钥加密数据
func EncryptRSA(publicKey *rsa.PublicKey, data []byte) ([]byte, error) {
	hash := sha256.New() // 使用 SHA-256 哈希
	ciphertext, err := rsa.EncryptOAEP(hash, rand.Reader, publicKey, data, nil)
	if err != nil {
		return nil, fmt.Errorf("加密失败: %v", err)
	}
	return ciphertext, nil
}

// DecryptRSA 使用私钥解密数据
func DecryptRSA(privateKey *rsa.PrivateKey, ciphertext []byte) ([]byte, error) {
	hash := sha256.New() // 使用 SHA-256 哈希
	plaintext, err := rsa.DecryptOAEP(hash, rand.Reader, privateKey, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("解密失败: %v", err)
	}
	return plaintext, nil
}

// ExportPrivateKey 将私钥导出为 PEM 格式
func ExportPrivateKey(privateKey *rsa.PrivateKey) ([]byte, error) {
	privKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	block := &pem.Block{Type: "PRIVATE KEY", Bytes: privKeyBytes}
	return pem.EncodeToMemory(block), nil
}

// ExportPublicKey 将公钥导出为 PEM 格式
func ExportPublicKey(publicKey *rsa.PublicKey) ([]byte, error) {
	pubKeyBytes := x509.MarshalPKCS1PublicKey(publicKey)
	block := &pem.Block{Type: "PUBLIC KEY", Bytes: pubKeyBytes}
	return pem.EncodeToMemory(block), nil
}
